# -*- coding: utf-8 -*-
"""
修正后的 LangChain 对话系统基础框架
注意：此代码为示例，展示了如何修正主要错误。实际应用需根据具体需求完善逻辑。
"""

import os
# 确保必要的库已安装: pip install langchain langchain-community langchain-ollama faiss-cpu langchain-huggingface

# 导入修正
# 注意版本兼容性和弃用警告
from datetime import datetime
# from langchain.chains.question_answering.map_rerank_prompt import output_parser # 通常不直接使用
from langchain.chains import SequentialChain, LLMChain # RouterChain 用法复杂，暂不使用
# from langchain_community.embeddings import OpenAIEmbeddings # 如果使用 OpenAI
from langchain_community.llms import OpenAI # 修正导入
from langchain_core.callbacks import StreamingStdOutCallbackHandler
from langchain_ollama import OllamaLLM
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
# from langchain.memory import Memory # 错误导入，Memory 是抽象概念
from langchain.memory import ConversationBufferMemory, ConversationSummaryMemory
import sqlite3

# --- 配置 ---
# 最大历史记录长度
MAX_HISTORY_LENGTH = 5
# 对话轮数计数器
dialogue_count = 0

# --- LLM 初始化 ---
# OpenAI 示例 (需要 API Key)
# llm_openai = OpenAI(
#     api_key=os.getenv("OPENAI_API_KEY", "your-default-key-if-env-not-set"),
#     model="gpt-3.5-turbo-instruct", # OpenAI LLM 示例
#     base_url="https://api.openai.com/v1/", # 正确的 OpenAI API URL
#     streaming=True,
#     callbacks=[StreamingStdOutCallbackHandler()]
# )

# Ollama 示例 (本地运行)
llm_ollama = OllamaLLM(
    model="qwen3:4b", # 确保该模型已通过 Ollama 下载
    temperature=0.3,
    num_ctx=4096,
    # stream=True, # 注意：在某些链中直接使用 stream 可能需要特殊处理
    callbacks=[StreamingStdOutCallbackHandler()]
)

# 选择要使用的 LLM
llm = llm_ollama # 或 llm_openai

# --- Embeddings 初始化 ---
# OpenAI Embeddings 示例
# embeddings_model = OpenAIEmbeddings(api_key=os.getenv("OPENAI_API_KEY"))

# HuggingFace Embeddings 示例 (需要模型文件)
# 确保 'model/bge-large-zh' 路径正确且模型已下载
embeddings_model = HuggingFaceEmbeddings(
    model_name="model/bge-large-zh", # 或使用 'BAAI/bge-large-zh-v1.5' 并在线加载
    model_kwargs={'device': 'cpu'} # 根据环境改为 'cuda' 或 'cpu'
)

# --- Memory 初始化 ---
# 创建聊天记忆缓存区 (缓冲区记忆)
buffer_memory = ConversationBufferMemory(memory_key="chat_history", input_key="user_input", output_key="response")

# 创建摘要缓存 (摘要记忆)
summary_memory = ConversationSummaryMemory(llm=llm, memory_key="chat_summary", input_key="user_input")
# 注意：直接设置 llm 属性可能不适用于所有版本，通常在初始化时传入
# summary_memory.llm.max_tokens = 100 # 控制摘要长度 (可能需要通过 llm_kwargs 或 prompt 控制)
# summary_memory.llm.temperature = 0.5 # 控制生成稳定性 (同上)

# --- Prompt Template 定义 ---
# 修正 PromptTemplate 的定义
# 一个简单的问答模板
qa_prompt_template = PromptTemplate.from_template("基于以下上下文回答问题:\n{context}\n\n问题: {question}\n答案:")

# 一个带有历史记录的对话模板
# 注意：这里的 {chat_history} 需要与 Memory 的 memory_key 匹配
conversation_prompt_template = PromptTemplate.from_template(
    "你是一个有帮助的助手。\n"
    "之前的对话历史:\n{chat_history}\n\n"
    "用户: {user_input}\n"
    "助手:"
)

# --- 向量数据库 ---
# 初始化向量数据库 (这里创建一个空的 FAISS 索引)
# 通常你会先加载或创建一个包含文档的 vectorstore
vector_store = FAISS.from_texts(["初始化文本"], embedding=embeddings_model) # 创建一个初始空库或用示例文本

# --- 核心函数 ---

# 向量化并存储对话 (简化版)
def store_conversation_as_vector(vector_store, user_input, model_response):
    try:
        # FAISS 的 add_texts 方法会自动处理嵌入
        texts = [f"用户: {user_input}", f"助手: {model_response}"]
        vector_store.add_texts(texts)
        print("对话已存储到向量数据库。")
    except Exception as e:
        print(f"存储对话时出错: {e}")

# 检索相关对话 (简化版)
def retrieve_relevant_conversation(vector_store, query, k=2):
    try:
        # FAISS 的 similarity_search 方法
        docs = vector_store.similarity_search(query, k=k)
        # 提取文本内容
        relevant_texts = [doc.page_content for doc in docs]
        print(f"检索到相关历史: {relevant_texts}")
        return relevant_texts
    except Exception as e:
        print(f"检索相关对话时出错: {e}")
        return []

# 动态调整记忆深度 (示例逻辑)
def dynamic_memory_update(memory, user_input, model_response, max_depth=5):
    global dialogue_count
    dialogue_count += 1
    # 简化处理：这里只演示保存上下文
    # 实际的 buffer_window 调整取决于具体 Memory 类型的实现
    # ConversationBufferMemory 通常通过 max_token_limit 控制
    memory.save_context({"user_input": user_input}, {"response": model_response})
    print(f"对话轮数: {dialogue_count}")

# 预处理输入
def preprocess(user_input):
    cleaned_input = user_input.strip()
    # 这里可以根据需要添加更多预处理逻辑
    return cleaned_input

# 查询数据库 (示例)
def query_database(keyword):
    try:
        # 确保 database.db 存在且有 topics 表
        conn = sqlite3.connect("database.db")
        cursor = conn.cursor()
        cursor.execute("SELECT info FROM topics WHERE name=?", (keyword,))
        result = cursor.fetchone()
        conn.close()
        return result[0] if result else '没有在数据库中找到相关信息'
    except Exception as e:
         print(f"查询数据库时出错: {e}")
         return "查询数据库失败"

# 简单路由判断
def determine_router(input_text):
    if "查询" in input_text:
        return "database_query"
    else:
        return "general_conversation"

# --- 主流程 ---
def main():
    print("启动对话系统... (输入 'exit' 退出)")
    while True:
        user_input = input("\n用户: ")
        if user_input.lower() == 'exit':
            print("再见！")
            break

        processed_input = preprocess(user_input)
        route = determine_router(processed_input)

        if route == "database_query":
            # 提取关键词进行查询 (这里简化为整个输入)
            keyword = processed_input.replace("查询", "").strip()
            db_result = query_database(keyword)
            print(f"\n助手: {db_result}")
            # 可以选择是否将数据库查询结果存入记忆或向量库
            dynamic_memory_update(buffer_memory, processed_input, db_result)

        else: # general_conversation
            # 1. 检索相关历史
            relevant_history = retrieve_relevant_conversation(vector_store, processed_input)

            # 2. 构建上下文 (这里简单拼接)
            context = "\n".join(relevant_history) if relevant_history else "无相关历史记录"

            # 3. 构建 Prompt (使用带历史的模板)
            # 注意：这里使用 buffer_memory 加载历史
            memory_vars = buffer_memory.load_memory_variables({})
            # print(f"Memory Vars: {memory_vars}") # 调试用
            full_prompt = conversation_prompt_template.format(
                chat_history=memory_vars.get('chat_history', ''),
                user_input=processed_input
            )
            print(f"\n构建的 Prompt:\n{full_prompt}\n") # 调试用

            # 4. 调用 LLM 生成回复
            try:
                # 使用 invoke 或 __call__ (取决于 LangChain 版本)
                response_obj = llm.invoke(full_prompt) # OllamaLLM 通常返回字符串
                # 如果使用 OpenAI 或其他返回对象的 LLM，可能需要 response_obj.choices[0].text.strip()
                model_response = response_obj.strip() if isinstance(response_obj, str) else str(response_obj)
                print(f"\n助手: {model_response}")

                # 5. 更新记忆
                dynamic_memory_update(buffer_memory, processed_input, model_response)

                # 6. 存储对话到向量库
                store_conversation_as_vector(vector_store, processed_input, model_response)

            except Exception as e:
                print(f"\n调用 LLM 时出错: {e}")
                print("助手: 抱歉，我无法生成回复。")

    # 查看最终记忆状态 (可选)
    print("\n--- 最终对话历史 (Buffer Memory) ---")
    print(buffer_memory.load_memory_variables({}))
    print("\n--- 最终对话摘要 (Summary Memory) ---")
    print(summary_memory.load_memory_variables({})) # 如果启用了摘要记忆

if __name__ == "__main__":
    main()



